#!/usr/bin/env python3
# Copyright (C) 2024-present Naver Corporation. All rights reserved.
# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
#
# --------------------------------------------------------
# Script to pre-process the scannet++ dataset.
# Usage:
# python3 datasets_preprocess/preprocess_scannetpp.py --scannetpp_dir /path/to/scannetpp --precomputed_pairs /path/to/scannetpp_pairs --pyopengl-platform egl
# --------------------------------------------------------
import os
os.environ["PYOPENGL_PLATFORM"] = "egl" 
import OpenGL
print("PyOpenGL platform:", os.environ.get("PYOPENGL_PLATFORM"))
print("Using OpenGL implementation at:", OpenGL.__file__)

import argparse
import os.path as osp
import re
from tqdm import tqdm
import json
from scipy.spatial.transform import Rotation
import pyrender
import trimesh
import trimesh.exchange.ply
import numpy as np
import cv2
import PIL.Image as Image
import sys
import shutil
import logging
import rootutils
rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)

# Setup logger
logger = logging.getLogger()
logger.setLevel(logging.INFO)  # Set minimum log level

# Console handler
console_handler = logging.StreamHandler()
console_handler.setLevel(logging.INFO)

# File handler
file_handler = logging.FileHandler('/home/liucong/data/3d/scannetpp_processed1/preprocess.log', mode='a')
file_handler.setLevel(logging.INFO)

# Formatter (without time)
formatter = logging.Formatter('%(asctime)s %(levelname)s: %(message)s')

# Add formatter to handlers
console_handler.setFormatter(formatter)
file_handler.setFormatter(formatter)

# Add handlers to logger
logger.addHandler(console_handler)
logger.addHandler(file_handler)

# current_dir = os.path.dirname(os.path.abspath(__file__))
# src_dir = os.path.abspath(os.path.join(current_dir, "..", "src"))
# sys.path.insert(0, src_dir)

from utils.cropping import rescale_image_depthmap
import src.dust3r.utils.geometry as geometry

inv = np.linalg.inv
norm = np.linalg.norm
# REGEXPR_DSLR = re.compile(r"^DSC(?P<frameid>\d+).JPG$")
REGEXPR_DSLR = re.compile(r"^(?:[0-9a-f]{8}_)?DSC(?P<frameid>\d+)\.JPG$")
REGEXPR_IPHONE = re.compile(r"frame_(?P<frameid>\d+).jpg$")

DEBUG_VIZ = None  # 'iou'
if DEBUG_VIZ is not None:
    import matplotlib.pyplot as plt  # noqa


OPENGL_TO_OPENCV = np.float32(
    [[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]]
)


def get_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument("--scannetpp_dir", type=str, default="/home/liucong/data/3d/scannetpp_unzipped")
    parser.add_argument("--mesh_dir", type=str, default=None)
    parser.add_argument("--output_dir", type=str, default="/home/liucong/data/3d/scannetpp_processed1")
    parser.add_argument(
        "--target_resolution", default=576, type=int, help="images resolution (576, 384)"
    )
    parser.add_argument(
    "--render_depth",
    dest="render_depth",
    action="store_true",
    help="Whether to render depth")
    
    return parser


def pose_from_qwxyz_txyz(elems):
    qw, qx, qy, qz, tx, ty, tz = map(float, elems)
    pose = np.eye(4)
    pose[:3, :3] = Rotation.from_quat((qx, qy, qz, qw)).as_matrix()
    pose[:3, 3] = (tx, ty, tz)
    return np.linalg.inv(pose)  # returns cam2world


def get_frame_number(name, cam_type="dslr"):
    if cam_type == "dslr":
        regex_expr = REGEXPR_DSLR
    elif cam_type == "iphone":
        regex_expr = REGEXPR_IPHONE
    else:
        raise NotImplementedError(f"wrong {cam_type=} for get_frame_number")
    if '_' in name:
        name = name.split('_')[1]
    matches = re.match(regex_expr, name)
    return matches["frameid"]


def load_sfm(sfm_dir, cam_type="dslr"):
    # load cameras
    with open(osp.join(sfm_dir, "cameras.txt"), "r") as f:
        raw = f.read().splitlines()[3:]  # skip header

    intrinsics = {}
    for camera in tqdm(raw, position=1, leave=False):
        camera = camera.split(" ")
        intrinsics[int(camera[0])] = [camera[1]] + [float(cam) for cam in camera[2:]]

    # load images
    with open(os.path.join(sfm_dir, "images.txt"), "r") as f:
        raw = f.read().splitlines()
        raw = [line for line in raw if not line.startswith("#")]  # skip header

    img_idx = {}
    img_infos = {}
    for image, points in tqdm(
        zip(raw[0::2], raw[1::2]), total=len(raw) // 2, position=1, leave=False
    ):
        image = image.split(" ")
        points = points.split(" ")

        idx = image[0]
        img_name = image[-1]
        assert img_name not in img_idx, "duplicate db image: " + img_name
        img_idx[img_name] = idx  # register image name

        current_points2D = {
            int(i): (float(x), float(y))
            for i, x, y in zip(points[2::3], points[0::3], points[1::3])
            if i != "-1"
        }
        img_infos[idx] = dict(
            intrinsics=intrinsics[int(image[-2])],
            path=img_name,
            frame_id=get_frame_number(img_name, cam_type),
            cam_to_world=pose_from_qwxyz_txyz(image[1:-2]),
            sparse_pts2d=current_points2D,
        )

    # load 3D points
    with open(os.path.join(sfm_dir, "points3D.txt"), "r") as f:
        raw = f.read().splitlines()
        raw = [line for line in raw if not line.startswith("#")]  # skip header

    points3D = {}
    observations = {idx: [] for idx in img_infos.keys()}
    for point in tqdm(raw, position=1, leave=False):
        point = point.split()
        point_3d_idx = int(point[0])
        points3D[point_3d_idx] = tuple(map(float, point[1:4]))
        if len(point) > 8:
            for idx, point_2d_idx in zip(point[8::2], point[9::2]):
                observations[idx].append((point_3d_idx, int(point_2d_idx)))

    return img_idx, img_infos, points3D, observations


def subsample_img_infos(img_infos, num_images, allowed_name_subset=None):
    img_infos_val = [(idx, val) for idx, val in img_infos.items()]
    if allowed_name_subset is not None:
        img_infos_val = [
            (idx, val)
            for idx, val in img_infos_val
            if val["path"] in allowed_name_subset
        ]

    if len(img_infos_val) > num_images:
        img_infos_val = sorted(img_infos_val, key=lambda x: x[1]["frame_id"])
        kept_idx = (
            np.round(np.linspace(0, len(img_infos_val) - 1, num_images))
            .astype(int)
            .tolist()
        )
        img_infos_val = [img_infos_val[idx] for idx in kept_idx]
    return {idx: val for idx, val in img_infos_val}


def undistort_images(intrinsics, rgb, mask):
    camera_type = intrinsics[0]

    width = int(intrinsics[1])
    height = int(intrinsics[2])
    fx = intrinsics[3]
    fy = intrinsics[4]
    cx = intrinsics[5]
    cy = intrinsics[6]
    distortion = np.array(intrinsics[7:])

    K = np.zeros([3, 3])
    K[0, 0] = fx
    K[0, 2] = cx
    K[1, 1] = fy
    K[1, 2] = cy
    K[2, 2] = 1

    K = geometry.colmap_to_opencv_intrinsics(K)
    if camera_type == "OPENCV_FISHEYE":
        assert len(distortion) == 4

        new_K = cv2.fisheye.estimateNewCameraMatrixForUndistortRectify(
            K,
            distortion,
            (width, height),
            np.eye(3),
            balance=0.0,
        )
        # Make the cx and cy to be the center of the image
        new_K[0, 2] = width / 2.0
        new_K[1, 2] = height / 2.0

        map1, map2 = cv2.fisheye.initUndistortRectifyMap(
            K, distortion, np.eye(3), new_K, (width, height), cv2.CV_32FC1
        )
    else:
        new_K, _ = cv2.getOptimalNewCameraMatrix(
            K, distortion, (width, height), 1, (width, height), True
        )
        map1, map2 = cv2.initUndistortRectifyMap(
            K, distortion, np.eye(3), new_K, (width, height), cv2.CV_32FC1
        )

    undistorted_image = cv2.remap(
        rgb,
        map1,
        map2,
        interpolation=cv2.INTER_LINEAR,
        borderMode=cv2.BORDER_REFLECT_101,
    )
    undistorted_mask = cv2.remap(
        mask,
        map1,
        map2,
        interpolation=cv2.INTER_LINEAR,
        borderMode=cv2.BORDER_CONSTANT,
        borderValue=255,
    )
    K = geometry.opencv_to_colmap_intrinsics(K)
    return width, height, new_K, undistorted_image, undistorted_mask
def process_scenes(scannetpp_dir, output_dir, mesh_dir=None,target_resolution=512, render_depth=True):
    """
    Process each scene in the Scannet++ dataset using DSLR images only.
    
    The unprocessed Scannet++ directory is structured as follows:
      scannetpp/
         data/
             <scene_name>/
                 dslr/
                     colmap/
                     nerfstudio/
                     resized_anon_masks/
                     resized_images/
                 scans/
                     mesh_aligned_0.05.ply
         metadata/
         splits/
    
    For each scene (located under "data"), this function:
      - Checks for the presence of a "dslr" folder and a mesh file in the "scans" folder.
      - Loads the DSLR SfM data (from dslr/colmap).
      - Enumerates all DSLR RGB images (files matching the pattern DSC*.JPG) in "dslr/resized_images".
      - For each image, loads the corresponding mask from "dslr/resized_anon_masks", undistorts the image and mask,
        and then rescales them to the target resolution (width = target_resolution, height = 3/4 * target_resolution).
      - Renders a depth map using Pyrender from the mesh in "scannetpp/data/<scene>/scans".
      - Saves, in the output directory, all images for that scene in a single folder with files:
            "xxxxxx_rgb.JPG"  (RGB image)
            "xxxxxx_depth.png" (depth map, in millimeters, as uint16)
            "xxxxxx_cam.npz"   (camera metadata containing intrinsics and pose)
    """
    # The scenes are stored under the "data" folder inside scannetpp_dir.
    data_dir = osp.join(scannetpp_dir, "data")
    # scenes = sorted([d for d in os.listdir(data_dir) if osp.isdir(osp.join(data_dir, d))])
    scenes = ['0271889ec0']
    
    # Depth rendering parameters.
    znear = 0.05
    zfar = 20.0

    if render_depth:
        import pyrender

    for scene in tqdm(scenes, desc="Processing scenes"):
        scene_dir = osp.join(data_dir, scene)
        # Check for the DSLR subdirectory.
        dslr_dir = osp.join(scene_dir, "dslr")
        if not osp.isdir(dslr_dir):
            logger.error(f"DSLR directory not found in scene '{scene}'. Skipping.")
            continue
        
        if render_depth:
            # Get the mesh file from the scans folder.
            if mesh_dir is None:
                scene_mesh_dir = scene_dir
            else:
                scene_mesh_dir = osp.join(mesh_dir, "data", scene)
            scans_dir = osp.join(scene_mesh_dir, "scans")
            ply_path = osp.join(scans_dir, "mesh_aligned_0.05.ply")
            if not osp.isfile(ply_path):
                logger.error(f"Mesh file not found in scene '{scene}'. Skipping.")
                print(f"Mesh file not found in scene '{scene}'. Skipping.")
                continue
            # Load the mesh using trimesh and add it to a Pyrender scene.
            
            try:
                with open(ply_path, "rb") as f:
                    mesh_kwargs = trimesh.exchange.ply.load_ply(f)  
                mesh_scene = trimesh.Trimesh(**mesh_kwargs)

                # Remove duplicate faces
                mesh_scene.update_faces(mesh_scene.unique_faces())
                # Remove degenerate faces (use a small height threshold, e.g., 1e-5)
                mesh_scene.update_faces(mesh_scene.nondegenerate_faces())
                mesh_scene.merge_vertices()

                if len(mesh_scene.faces) > 10000000:
                    target_faces = 10000000
                    # mesh_scene.faces.shape[0] // 2 
                    mesh_scene = mesh_scene.simplify_quadric_decimation(face_count=target_faces)

                mesh = pyrender.Mesh.from_trimesh(mesh_scene, smooth=False)
                pyrender_scene = pyrender.Scene()
                pyrender_scene.add(mesh)
                # Create a Pyrender offscreen renderer (viewport dimensions will be updated per image).
                renderer = pyrender.OffscreenRenderer(1, 1)

            except Exception as e:
                logger.error(f"Error loading mesh: {e} in {scene}")
                return False
        
        # Load DSLR SfM data.
        sfm_dir = osp.join(dslr_dir, "colmap")
        rgb_dir = osp.join(dslr_dir, "resized_images")
        mask_dir = osp.join(dslr_dir, "resized_anon_masks")
        # load_sfm returns (img_idx, img_infos, points3D, observations)
        img_idx, img_infos, _, _ = load_sfm(sfm_dir, cam_type="dslr")
        
        # Enumerate all DSLR RGB images (matching DSC*.JPG) in the "resized_images" folder.
        rgb_files = sorted([f for f in os.listdir(rgb_dir) if f.endswith(".JPG")])
        if len(rgb_files) == 0:
            logger.error(f"No DSLR images found in '{rgb_dir}'. Skipping scene '{scene}'.")
            continue
        
        # Create the output directory for this scene.
        scene_outdir = osp.join(output_dir, scene)
        os.makedirs(scene_outdir, exist_ok=True)

        fail = False
        for rgb_filename in tqdm(rgb_files, desc=f"Processing DSLR images in {scene}"):
            # Retrieve SfM metadata for the current image.
            if rgb_filename not in img_idx:
                logger.error(f"Image {rgb_filename} in {scene} not found in SfM data. Skipping.")
                fail = True
                continue
            idx = img_idx[rgb_filename]
            img_info = img_infos[idx]
            
            frame_id = img_info["frame_id"]
            
            # Build paths for the RGB image and its corresponding mask.
            rgb_path = osp.join(rgb_dir, rgb_filename)
            mask_filename = rgb_filename[:-3] + "png"
            mask_path = osp.join(mask_dir, mask_filename)
            if not osp.isfile(rgb_path) or not osp.isfile(mask_path):
                logger.error(f"RGB or mask file missing for image {rgb_filename} in {scene}")
                fail = True
                continue
                        # Build output filenames using six-digit frame id.
            out_rgb = osp.join(scene_outdir, frame_id + "_rgb.jpg")
            out_depth = osp.join(scene_outdir, frame_id + "_depth.png")
            out_cam = osp.join(scene_outdir, frame_id + "_cam.npz")

            # Load the images.
            rgb_img = np.array(Image.open(rgb_path))
            mask_img = np.array(Image.open(mask_path))
            
            # Undistort image and mask.
            width, height, new_K, undistorted_rgb, undistorted_mask = undistort_images(
                img_info["intrinsics"], rgb_img, mask_img
            )

            # rescale_image_depthmap assumes opencv intrinsics
            intrinsics = geometry.colmap_to_opencv_intrinsics(new_K)

            # Rescale to target resolution:
            width = target_resolution
            height = int(target_resolution * 2 / 3)
            # target_w = target_resolution
            # target_h = target_resolution
            undistorted_rgb_rescaled, undistorted_mask_rescaled, intrinsics = rescale_image_depthmap(
                undistorted_rgb, undistorted_mask, intrinsics, (width, height)
            )
            
            W, H = undistorted_rgb_rescaled.size
            intrinsics = geometry.opencv_to_colmap_intrinsics(intrinsics)
            print(W, H)

            undistorted_rgb_rescaled.save(out_rgb, quality=95)
            np.savez(out_cam, intrinsics=intrinsics, pose=img_info["cam_to_world"])

            if render_depth:
                # Update renderer viewport.
                renderer.viewport_width, renderer.viewport_height = W, H
                
                # Set up camera intrinsics.
                fx, fy, cx, cy = (
                    intrinsics[0, 0],
                    intrinsics[1, 1],
                    intrinsics[0, 2],
                    intrinsics[1, 2],
                )
                cam = pyrender.camera.IntrinsicsCamera(fx, fy, cx, cy, znear=znear, zfar=zfar)
                
                # Compute the camera-to-world pose.
                camera_node = pyrender_scene.add(
                    cam, pose=img_info["cam_to_world"] @ OPENGL_TO_OPENCV
                )
                # Render the depth map.
                depth = renderer.render(pyrender_scene, flags=pyrender.RenderFlags.DEPTH_ONLY)
                pyrender_scene.remove_node(camera_node)
                
                # Convert depth from meters to millimeters and cast to uint16.
                depth_mm = (depth * 1000).astype(np.uint16)
                # Invalidate depth where the rescaled mask is not valid.
                depth_mm[undistorted_mask_rescaled < 255] = 0
                
                # Save outputs.
                Image.fromarray(depth_mm).save(out_depth)
        
        if render_depth:
            # Clean up the renderer.
            renderer.delete()
            del pyrender_scene
        
        if fail:
            logger.error("Failed!{scene}")
        else:
            # shutil.rmtree(scene_dir)
            logger.info(f"Deleted {scene}")

        # break
    print("All scenes processed.")

if __name__ == "__main__":
    parser = get_parser()
    args = parser.parse_args()    
    # if args.pyopengl_platform.strip():
    #     os.environ["PYOPENGL_PLATFORM"] = args.pyopengl_platform
    # print(f"Scannet++ directory: {args.scannetpp_dir}")
    # print(f"Mesh directory: {args.mesh_dir}")
    # print(f"Output directory: {args.output_dir}")
    # print(f"Target resolution: {args.target_resolution}")
    # print(f"Render depth: {args.render_depth}")
    process_scenes(
        args.scannetpp_dir,
        args.output_dir,
        args.mesh_dir,
        args.target_resolution,
        args.render_depth,
    )